% Crank-Nicolson 法求解二维Schrodinger方程
% i \hbar \pdv{u}{t} = -\frac{\hbar^2}{2m} \pdv[2]{u}{x} + V u
% \frac{u^{(k+1)}_i - u^{(k)}_i}{\Delta t} = i\frac{\hbar}{2m} \frac{1}{2} \left( \frac{u^{(k+1)}_{i+1} - 2u^{(k+1)}_i + u^{(k+1)}_{i-1}}{(\Delta x)^2} + \frac{u^{(k)}_{i+1} - 2u^{(k)}_i + u^{(k)}_{i-1}}{(\Delta x)^2} \right)- \frac{i}{\hbar} V \frac{u^{(k+1)} + u^{(k)}}{2}
% u^{(k+1)}_i - \frac{\alpha}{2} \left( u^{(k+1)}_{i+1} - 2u^{(k+1)}_i + u^{(k+1)}_{i-1} \right) + \frac{\beta}{2} V u^{(k+1)}_i =  u^{(k)}_i + \frac{\alpha}{2} \left( u^{(k)}_{i+1} - 2u^{(k)}_i + u^{(k)}_{i-1} \right) - \frac{\beta}{2} V u^{(k)}_i
% 其中 \alpha = i\frac{\hbar}{2m} \cdot \frac{\Delta t}{(\Delta x)^2}, \quad \beta = \frac{i}{\hbar} \Delta t
% A u^{(k+1)} = B u^{(k)}, \quad
% A = I - \frac{\alpha}{2} \nabla^2 + \frac{\beta}{2} V, \quad
% B = I + \frac{\alpha}{2} \nabla^2 - \frac{\beta}{2} V
% 参考：
% https://wuli.wiki/online/CraNic.html
% https://zhuanlan.zhihu.com/p/393374195
% https://www.bilibili.com/video/BV1sq4y1V7JT
% 可能有bug
% Gitee Repo

clc
clear

global n

L = 2;
dx = 0.05;
dt = 0.01;

hbar = 1;
m = 5; % 质量
v = 4;
p = m*v;

[y,x] = meshgrid(-L:dx:L);
n = size(x,1);

% 初始波包
u0 = exp(-20*((x+0.5*L).^2+y.^2)).*exp(i*p/hbar*x);
u0(1,:)=0;
u0(n,:)=0;
u0(:,1)=0;
u0(:,n)=0;

_u0 = u0(:);

% 势能
V = zeros(n,n);
# 双缝
V(abs(x)<0.04*L) = 100;
V(abs(y-0.1*L)<0.05*L | abs(y+0.1*L)<0.05*L) = 0;

# 单侧高能垒
##V(x>0)=100;

alpha = i*hbar/(2*m)*dt/(dx)^2;
beta = i/hbar*dt;

function A = compute_laplacian_old()
    global n
    A = spalloc(n^2, n^2, 5*n^2)
    for i = 1:n
      for j = 1:n
        _k = i+(j-1)*n;

        if i == 1 || i == n || j == 1 || j == n %如果是边界
          A(_k,_k) = 1;
        else
          A(_k,_k) = -4;
          A(_k,_k+1) = 1;
          A(_k,_k-1) = 1;
          A(_k,_k+n) = 1;
          A(_k,_k-n) = 1;
        end
      end
    end
    return;

end

function lap = compute_laplacian()
    global n

    lap = spalloc(n*n,n*n);

    [j,i] = meshgrid(2:n-1);
    i = i(:);j = j(:);
    ind = zeros((n-2)^2,1);
    ind(:,1) = sub2ind([n,n],i,j);
    ind(:,2) = sub2ind([n,n],i+1,j);
    ind(:,3) = sub2ind([n,n],i-1,j);
    ind(:,4) = sub2ind([n,n],i,j+1);
    ind(:,5) = sub2ind([n,n],i,j-1);

    ind2 = sub2ind([n^2,n^2],ind(:,1),ind(:,1));
    lap(ind2)= -4;

    for k = 2:5
        ind2 = sub2ind([n^2,n^2],ind(:,1),ind(:,k));
        lap(ind2) = 1;
    end

    i = (1:n)';
    k0 = 1+0*i;
    k1 = n+0*i;

    ind = sub2ind([n,n],k0,i);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    lap(ind2) = 1;

    ind = sub2ind([n,n],k1,i);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    lap(ind2) = 1;

    ind = sub2ind([n,n],i,k0);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    lap(ind2) = 1;

    ind = sub2ind([n,n],i,k1);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    lap(ind2) = 1;
end

function Vmat = compute_Vmat(V)
    global n

    Vmat = spalloc(n*n,n*n);
    [j,i] = meshgrid(2:n-1);
    i = i(:);j = j(:);
    ind = sub2ind([n,n],i,j);
    ind2 = sub2ind([n^2,n^2],ind,ind);
    Vmat(ind2) = V(ind);
    return;
end

lap = compute_laplacian();
Vmat = compute_Vmat(V);

clear i

A = eye(n^2) - 1/2*alpha*lap + 1/2*beta*Vmat;
B = eye(n^2) + 1/2*alpha*lap - 1/2*beta*Vmat;

for tick = 1:2000
    _u1 = A\(B*_u0);
    _u0 = _u1;

    if mod(tick,10) == 0
        clf
        hold on
        axis equal
        axis([-L,L,-L,L,0,1])

        u1 = reshape(_u1,n,n);
        umag = real(u1).^2+imag(u1).^2;

        surf(x,y,umag,'edgecolor','none')

        xlabel('x');
        ylabel('y');
        zlabel('Psi');
        view([30,45])

        drawnow
        pause(0.1)

        disp(sum(umag(:)))
    end
end

